package com.zcs.shardingjdbc.algorithm;

import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingValue;
import org.springframework.core.env.Environment;
import org.springframework.util.StringUtils;

import java.sql.*;
import java.time.YearMonth;
import java.time.format.DateTimeFormatter;
import java.time.format.DateTimeParseException;
import java.util.*;

/**
 * @author Houger.
 * @date 2024年08月16日 17:55.
 */
@Slf4j
public class TimeShardingAlgorithm implements PreciseShardingAlgorithm<String> {



    private static final HashSet<String> tableNameCache = new HashSet<>();

    /** 数据库配置 */
    private static final Environment ENV = SpringUtil.getApplicationContext().getEnvironment();
    private static final String DATASOURCE_URL = ENV.getProperty("spring.shardingsphere.datasource.o0.url");
    private static final String DATASOURCE_USERNAME = ENV.getProperty("spring.shardingsphere.datasource.o0.username");
    private static final String DATASOURCE_PASSWORD = ENV.getProperty("spring.shardingsphere.datasource.o0.password");


    static {
        //读取数据库中所有表名
        List<String> tableList = getAllTableName();
        tableNameCache.addAll(tableList);

    }
    /**
     * 获取数据库中的表名
     */
    public static List<String> getAllTableName() {
        List<String> res = new ArrayList<>();
        // 获取数据中的表名
        try {
            assert DATASOURCE_URL != null;
            try (Connection connection = DriverManager.getConnection(DATASOURCE_URL, DATASOURCE_USERNAME, DATASOURCE_PASSWORD);
                     Statement st = connection.createStatement()) {
                try (ResultSet rs = st.executeQuery("SHOW TABLES")) {
                    while (rs.next()) {
                        res.add(rs.getString(1));
                    }
                }
            }
        } catch (Exception e) {
            log.error(e.getMessage());
            throw new IllegalArgumentException("SHOW TABLES失败，请稍后重试");
        }
        return res;
    }



    /**
     * 
     * @author Mr.hao
     * @date 2024/8/16 17:59
     * @param availableTargetNames 可用的数据源或表的名称
     * @param shardingValue        PreciseShardingValue: 其中logicTableName为逻辑表，columnName分片键（字段），value为分片键的值
     * @return java.lang.String
     */
    @Override
    public String doSharding(final Collection<String> availableTargetNames, final PreciseShardingValue<String> shardingValue) {
        log.info("execute TimeShardingAlgorithm.doSharding....");
        log.info("【TimeShardingAlgorithm】availableTargetNames:{}, shardingValue:{}", availableTargetNames, shardingValue);
        String shardingIndex = shardingValue.getValue().substring(0,6);
        if (!isValidYearMonth(shardingIndex)) {
            throw new RuntimeException("分片不合法，不是合法的年月,shardingIndex="+shardingIndex);
        }

        synchronized (shardingValue.getLogicTableName().intern()) {
            // 缓存中有此表 返回
            if (tableNameCache.contains(shardingValue.getLogicTableName()+"_"+shardingIndex)) {
                return shardingValue.getLogicTableName()+"_"+shardingIndex;
            }
            // 缓存中无此表 建表 并添加缓存
            boolean res=createShardingTable(shardingValue.getLogicTableName(),shardingValue.getLogicTableName()+"_"+shardingIndex);
            if (res){
                tableNameCache.add(shardingValue.getLogicTableName()+"_"+shardingIndex);
            }
        }
        return shardingValue.getLogicTableName()+"_"+shardingIndex;
    }


    /**
     * 创建分表
     * @param logicTable 逻辑表
     * @param resultTableName 真实表名
     * @return 创建结果（true创建成功，false未创建）
     */
    private static boolean createShardingTable(String logicTable, String resultTableName) {
        // 根据日期判断，当前月份之后分表不提前创建
        String month = resultTableName.replace(logicTable + "_","");
        YearMonth shardingMonth = YearMonth.parse(month, DateTimeFormatter.ofPattern("yyyyMM"));
        if (shardingMonth.isAfter(YearMonth.now())) {
            throw new RuntimeException("分片时间过早 | Premature fragmentation time");
        }

        synchronized (logicTable.intern()) {
            executeSql(Collections.singletonList("CREATE TABLE IF NOT EXISTS `" + resultTableName + "` LIKE `" + logicTable + "`;"));
        }
        return true;
    }
    /**
     * 执行SQL
     * @param sqlList SQL集合
     */
    private static void executeSql(List<String> sqlList) {
        if (StringUtils.isEmpty(DATASOURCE_URL) || StringUtils.isEmpty(DATASOURCE_USERNAME) || StringUtils.isEmpty(DATASOURCE_PASSWORD)) {
            log.error(">>>>>>>>>> 【ERROR】数据库连接配置有误，请稍后重试，URL:{}, username:{}, password:{}", DATASOURCE_URL, DATASOURCE_USERNAME, DATASOURCE_PASSWORD);
            throw new IllegalArgumentException("数据库连接配置有误，请稍后重试");
        }
        try (Connection conn = DriverManager.getConnection(DATASOURCE_URL, DATASOURCE_USERNAME, DATASOURCE_PASSWORD)) {
            try (Statement st = conn.createStatement()) {
                conn.setAutoCommit(false);
                for (String sql : sqlList) {
                    st.execute(sql);
                }
            } catch (Exception e) {
                conn.rollback();
                log.error(">>>>>>>>>> 【ERROR】数据表创建执行失败，请稍后重试，原因：{}", e.getMessage(), e);
                throw new IllegalArgumentException("数据表创建执行失败，请稍后重试");
            }
        } catch (SQLException e) {
            log.error(">>>>>>>>>> 【ERROR】数据库连接失败，请稍后重试，原因：{}", e.getMessage(), e);
            throw new IllegalArgumentException("数据库连接失败，请稍后重试");
        }
    }

    public static boolean isValidYearMonth(String str) {
        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyyMM");
        try {
            YearMonth.parse(str, formatter);
            return true;
        } catch (DateTimeParseException e) {
            return false;
        }
    }
}
